# 对文件加密解密操作


import hashlib
import os
from Crypto.Cipher import AES


def encryptFile(file_path, password='my_password', iv='myIV' * 4, delete_old_file=False, file_suffix='encx'):
    """
    对文件加密并修改后缀名
    :param file_path: 加密文件的路径
    :param password: 加密密码 默认 my_password
    :param iv: iv默认 myIVmyIVmyIVmyIV
    :param delete_old_file: 是否删除加密前的文件
    :param file_suffix: 加密后文件后缀名 默认：encx
    :return: 加密后的文件路径
    """

    # 填充以便使用AES-128、AES-192或AES-256加密算法中的任何一个
    padding = '{'
    # 将密码散列处理
    key = hashlib.sha256(padding.encode() + password.encode()).digest()
    # iv = 'myIV'*4  # iv必须是16位长度
    # 初始化加密器
    cipher = AES.new(key, AES.MODE_CBC, IV=iv.encode())
    new_file_path = file_path + '.' + file_suffix
    # 填充并加密文件
    chunk_size = 64 * 1024  # 每次加密64 kb
    with open(file_path, 'rb') as f:
        with open(new_file_path, 'wb') as out_file:
            while True:
                chunk = f.read(chunk_size)
                if len(chunk) == 0:
                    break
                elif len(chunk) % 16 != 0:
                    chunk += ((16 - len(chunk) % 16) * padding).encode()
                out_file.write(cipher.encrypt(chunk))

    if delete_old_file:
        # 删除原始文件
        os.remove(file_path)

    return new_file_path


def decryptFile(file_path, password='my_password', iv='myIV' * 4, delete_old_file=False):
    """
    解密文件
    :param file_path: 加密文件的路径
    :param password: 加密密码 默认 my_password
    :param iv: iv默认 myIVmyIVmyIVmyIV
    :param delete_old_file: 是否删除加密前的文件
    :return: 解密后文件路径
    """
    # 填充
    padding = '{'
    # 密码散列处理
    key = hashlib.sha256(padding.encode() + password.encode()).digest()
    # 初始化解密器
    cipher = AES.new(key, AES.MODE_CBC, IV=iv.encode())
    # 解密文件
    chunk_size = 64 * 1024  # 64 kb
    with open(file_path, 'rb') as f:
        new_file_path = '.'.join(file_path.split('.')[:-1])
        with open(new_file_path, 'wb') as out_file:
            while True:
                chunk = f.read(chunk_size)
                if len(chunk) == 0:
                    break
                out_file.write(cipher.decrypt(chunk))

    if delete_old_file:
        # 删除加密文件
        os.remove(file_path)
    return new_file_path

path = "./Snipaste_2023-04-06_17-49-13.png"
# encryptFile(path)

decryptFile(path+".encx")